# -*- coding: utf-8 -*-
import asyncio
from asyncio.events import AbstractEventLoop
import random
import base64

import config
from db_tools import async_redis_cli
from logger import logger

INCREASE = 1
DECREASE = -1


class UpstreamProxyHandle:
    def __init__(self, loop):
        self.__redis_key = config.get('application', 'proxy_redis_key')
        self.loop: AbstractEventLoop = loop
        self.pu = config.get('application', 'proxy_user')
        self.pp = config.get('application', 'proxy_password')

    def __format_upstream(self, upstream_server):
        ip, port = upstream_server.split(':')
        return ip, int(port)

    async def __copy_data_thread(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
        try:
            while True:
                chunk = await reader.read(65536)
                if not chunk:
                    break
                writer.write(chunk)
                await writer.drain()
            writer.close()
        except Exception as e:
            return e
        else:
            return None

    async def __get_proxy_from_redis(self):
        results = await async_redis_cli.zrangebyscore(self.__redis_key, 0, 16, score_cast_func=int, withscores=True)
        if results:
            ip, score = random.choice(results)
            await self.__update_score(ip, score, INCREASE)
            __remote_ip, __remote_port = self.__format_upstream(ip)
            return __remote_ip, __remote_port, score

    @property
    def auth(self):
        auth = base64.b64encode(f"{self.pu}:{self.pp}".encode())
        return b'Proxy-Authorization: Basic ' + auth

    async def __update_score(self, key, score, type):
        assert type in [1, -1]
        while not await async_redis_cli.set(key, 1, nx=True, ex=60):
            # 获取锁
            await asyncio.sleep(1)
        await async_redis_cli.delete(key)  # 释放锁
        await async_redis_cli.zadd(self.__redis_key, {key: score + type}, nx=True)

    async def interject_reader(self, reader: asyncio.StreamReader):
        data = await reader.read(65536)
        if b'Proxy-Authorization' not in data and self.pu and self.pp:
            results = data.strip().split(b'\r\n')
            results.append(self.auth)
            data = b"\r\n".join(results) + b"\r\n\r\n"
        reader.feed_data(data)
        return reader

    @logger.catch
    async def download(self, a_reader: asyncio.StreamReader, a_writer: asyncio.StreamWriter):
        __remote_ip, __remote_port, score = await self.__get_proxy_from_redis()

        a_reader = await self.interject_reader(a_reader)

        logger.info(f'connect upstream server {__remote_ip}:{__remote_port}')

        try:
            c_reader, c_writer = await asyncio.open_connection(
                host=__remote_ip,
                port=__remote_port
            )
        except OSError:
            logger.warning('error open_connection')
        else:
            task_1 = asyncio.create_task(self.__copy_data_thread(c_reader, a_writer))
            task_2 = asyncio.create_task(self.__copy_data_thread(a_reader, c_writer))
            download_task = {task_1, task_2}
            await asyncio.wait(download_task, timeout=30, return_when=asyncio.FIRST_COMPLETED)
        finally:
            await self.__update_score(f"{__remote_ip}:{__remote_port}", score, DECREASE)


if __name__ == '__main__':
    UpstreamProxyHandle(loop=None)
